# -*- coding: utf-8 -*-
# @Time : 2022/3/28 18:16 -2022/03/30 12:10
# @Author : XyZeng

import os
import cv2
import numpy as np

import rgb2gray2rgb
import label2color

def label_color_label(img_path, color_map, debug=1):
    src_label = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)

    # resize
    # src_label = cv2.resize(img, (256, 256))
    '''

    label2color
    '''
    color_img = rgb2gray2rgb.gray2color(src_label, color_map)
    #save
    dir, name = os.path.split(img_path)
    save_path=os.path.join('./out_color', name)
    cv2.imwrite(save_path, color_img)

    '''
    color2label
    '''
    gray_img = rgb2gray2rgb.color2gray(color_img, color_map)

    '''
        原始lable图像的灰度值
        lable2color2label图像的灰度值,
    
    '''
    pix_value = np.unique(src_label)
    print('src_label_value', pix_value)
    pix_value = np.unique(gray_img)
    print("color2label_value", pix_value)

    if debug:
        cv2.imshow('label2color', color_img)
        cv2.imshow('color2label', gray_img)
        cv2.waitKey(0)


if __name__ == '__main__':

    img_dir = './label/'

    n_label = 20 # the  number of label class

    color_map1 = label2color.label_colormap(n_label=n_label, rgb1_bgr0=0)

    '''
    random gene color map
    '''
    random_state = np.random.RandomState(seed=1234)
    color_custom_map = np.random.randint(low=0, high=255, size=(n_label, 3), dtype=np.uint8)

    for img in sorted(os.listdir(img_dir)):

        if not img.endswith((".png", ".jpg")):
            continue

        img_path = os.path.join(img_dir, img)
        label_color_label(img_path, color_map1, debug=1)


# Ctrl + Alt + L     pycharm代码整理